import re
from typing import Dict, TypedDict
from langgraph.graph import END, StateGraph


# 定义状态结构
class AgentState(TypedDict):
    user_input: str  # 用户输入的字符串
    a: int           # 解析出的第一个数字
    b: int           # 解析出的第二个数字
    operation: str   # 解析出的操作类型
    result: int      # 计算结果


# 模拟大模型对话接口 - 解析用户输入
def mock_llm_chat(state: AgentState) -> Dict:
    """模拟大模型解析用户输入，提取数字和操作类型"""
    print(f"\n🧠 用户输入: '{state['user_input']}'")

    # 使用正则表达式提取数字
    numbers = re.findall(r'-?\d+', state["user_input"])

    # 简单规则判断操作类型
    if "加" in state["user_input"] or "+" in state["user_input"] or "和" in state["user_input"]:
        operation = "add"
    elif "减" in state["user_input"] or "-" in state["user_input"] or "差" in state["user_input"]:
        operation = "subtract"
    else:
        # 默认加法
        operation = "add"

    # 确保提取到两个数字
    if len(numbers) < 2:
        # 模拟大模型尝试修复
        print("⚠️ 模拟大模型: 未提取到足够数字，尝试补充默认值")
        numbers += ["0", "0"]  # 补充默认值
        numbers = numbers[:2]  # 只取前两个

    # 转换为整数
    try:
        a = int(numbers[0])
        b = int(numbers[1])
    except ValueError:
        a, b = 0, 0

    print(f"🔍 模拟大模型解析: 数字 {a} 和 {b}, 操作: {operation}")

    return {"a": a, "b": b, "operation": operation}


# 计算节点函数
def add_numbers(state: AgentState) -> Dict[str, int]:
    result = state["a"] + state["b"]
    print(f"➕ 计算结果: {state['a']} + {state['b']} = {result}")
    return {"result": result}


def subtract_numbers(state: AgentState) -> Dict[str, int]:
    result = state["a"] - state["b"]
    print(f"➖ 计算结果: {state['a']} - {state['b']} = {result}")
    return {"result": result}


# 结果展示节点
def display_result(state: AgentState) -> Dict:
    print(f"✅ 最终结果: {state['result']}")
    return {}


# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("chat", mock_llm_chat)
    workflow.add_node("add", add_numbers)
    workflow.add_node("subtract", subtract_numbers)
    workflow.add_node("display", display_result)

    # 设置入口点
    workflow.set_entry_point("chat")

    # 根据解析结果路由
    workflow.add_conditional_edges(
        "chat",
        lambda state: state["operation"],
        {
            "add": "add",
            "subtract": "subtract",
        }
    )

    # 连接计算节点到显示节点
    workflow.add_edge("add", "display")
    workflow.add_edge("subtract", "display")

    # 连接到结束
    workflow.add_edge("display", END)

    return workflow.compile()


# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 测试案例
    test_cases = [
        "5加3等于多少？",
        "10减去7",
        "计算20和15的和",
        "100 - 50",
        "8和4的差是多少？",
        "简单的计算：3和5",  # 测试默认操作
        "只有数字：42 100",  # 测试无操作符
    ]

    for input_str in test_cases:
        print("\n" + "="*50)
        print(f"测试输入: '{input_str}'")
        result = app.invoke({"user_input": input_str})
        print(f"完整输出: {result}")
